import os
import torchvision
from typing import Any, Callable, Dict, List, Optional, Tuple
from .yrb_dataset import YRBASIN

# import torch.nn as nn
# import torch.nn.functional as F
# Fu.relu()是函数调用，一般使用在foreward函数里。
# nn.ReLU()是模块调用，一般在定义网络层的时候使用。

# ROOT = f"e:\sfxData\YRBASIN\XTrain"
KINDS = ['x0', 'x1', 'x2', 'x3', 'x4']

YRBASIN_TRANSFORM_DEFAULT = torchvision.transforms.Compose([
    torchvision.transforms.Normalize((0.1307,), (0.3081,)),  # 设定标准化值
])


def get_yrbasin(root: str, data_flag: int = 0, transform: Optional[Callable] = None) -> YRBASIN:
    """
    取得 CIFAR10 数据集.

    参数:
        root: 数据集根目录
        train: 是否取得训练集
        transform: 转换
    """
    xtransform = YRBASIN_TRANSFORM_DEFAULT if transform == None else transform
    dset = YRBASIN(root=root, data_flag=data_flag, transform=xtransform)
    return dset


def get_yrbasin_kind_name(index: int):
    return KINDS[index]


def get_yrbasin_labels(labels) -> list[str]:
    return [KINDS[int(i)] for i in labels]

    # def _load_data(self):
